import os
import logging
from openai import OpenAI
import tkinter as tk
from tkinter import scrolledtext, messagebox
import time
import tkinter.font as tkFont
import docx
import hashlib
import pickle
import re
import jieba
from collections import defaultdict

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('investment_qa.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# 配置
DOCX_PATH = "D:/现代帕里斯帝用户手册/我的投资哲学.docx"
MAX_CONTEXT_LENGTH = 12000
MAX_RETRY = 3
RETRY_DELAY = 1
CACHE_FILE = "investment_cache.pkl"


class InvestmentQAApp:
    def __init__(self, root):
        self.root = root
        self.root.title("投资哲学咨询系统")
        self.root.geometry("1000x800")

        # 初始化jieba分词
        jieba.initialize()

        # 初始化字体
        self.available_fonts = list(tkFont.families())
        self.chinese_font = self.select_chinese_font()
        logger.info(f"使用字体: {self.chinese_font}")

        # 初始化豆包客户端
        self.client = OpenAI(
            base_url="https://ark.cn-beijing.volces.com/api/v3",
            api_key=os.environ.get("DOUBAO_API_KEY"),
        )

        # 加载文档
        self.doc_text = self.load_docx()
        if not self.doc_text:
            messagebox.showerror("错误", "无法加载投资哲学文档！")
            self.root.destroy()
            return

        # 预处理文档结构
        self.section_index = self.build_section_index()

        # 创建UI
        self.create_widgets()

    def select_chinese_font(self):
        """选择中文字体"""
        for font in ["Microsoft YaHei", "微软雅黑", "SimHei", "Arial"]:
            if font in self.available_fonts:
                return font
        return "TkDefaultFont"

    def load_docx(self):
        """加载Word文档"""
        try:
            # 缓存验证
            file_hash = hashlib.md5(open(DOCX_PATH, 'rb').read()).hexdigest()
            if os.path.exists(CACHE_FILE):
                with open(CACHE_FILE, 'rb') as f:
                    cache = pickle.load(f)
                    if cache.get('hash') == file_hash:
                        logger.info("从缓存加载文档")
                        return cache['text']

            # 解析Word文档
            logger.info(f"开始解析Word文档: {DOCX_PATH}")
            doc = docx.Document(DOCX_PATH)
            full_text = ""

            for para in doc.paragraphs:
                text = para.text.strip()
                if text:
                    if para.style.name.startswith(('Heading', '标题')):
                        full_text += f"\n【章节】{text}\n"
                    else:
                        full_text += text + "\n"

            # 保存缓存
            with open(CACHE_FILE, 'wb') as f:
                pickle.dump({'hash': file_hash, 'text': full_text}, f)

            logger.info(f"文档加载完成，共{len(full_text)}字符")
            return full_text

        except Exception as e:
            logger.error(f"文档加载失败: {str(e)}")
            return None

    def build_section_index(self):
        """构建章节索引（修复版）"""
        index = defaultdict(str)  # 使用defaultdict避免KeyError
        current_section = "前言"  # 默认章节名

        for line in self.doc_text.split('\n'):
            line = line.strip()
            if not line:
                continue

            if line.startswith('【章节】'):
                current_section = line[4:].strip()  # 移除【章节】前缀
                index[current_section] = ""
            else:
                index[current_section] += line + "\n"

        logger.info(f"已建立{len(index)}个章节索引")
        return dict(index)  # 转为普通字典返回

    def create_widgets(self):
        """创建界面"""
        # 标题
        title_frame = tk.Frame(self.root, bg="#2c3e50")
        title_frame.pack(fill=tk.X, padx=10, pady=10)

        title_label = tk.Label(
            title_frame,
            text="投资哲学咨询系统",
            font=(self.chinese_font, 18, "bold"),
            fg="white",
            bg="#2c3e50"
        )
        title_label.pack(pady=5)

        # 问题输入区
        input_frame = tk.Frame(self.root)
        input_frame.pack(fill=tk.X, padx=15, pady=10)

        tk.Label(
            input_frame,
            text="请输入您的投资问题:",
            font=(self.chinese_font, 12),
            anchor="w"
        ).pack(fill=tk.X)

        self.question_entry = tk.Entry(
            input_frame,
            font=(self.chinese_font, 14),
            width=80
        )
        self.question_entry.pack(fill=tk.X, pady=5)
        self.question_entry.bind("<Return>", lambda e: self.ask_question())

        # 按钮区
        button_frame = tk.Frame(self.root)
        button_frame.pack(pady=10)

        buttons = [
            ("提问", self.ask_question, "#3498db"),
            ("清空对话", self.clear_conversation, "#e74c3c"),
            ("退出系统", self.root.destroy, "#95a5a6")
        ]

        for text, command, color in buttons:
            btn = tk.Button(
                button_frame,
                text=text,
                command=command,
                font=(self.chinese_font, 12),
                bg=color,
                fg="white",
                width=15,
                relief=tk.FLAT
            )
            btn.pack(side=tk.LEFT, padx=10)

        # 对话历史区
        history_frame = tk.LabelFrame(
            self.root,
            text="咨询记录",
            font=(self.chinese_font, 12)
        )
        history_frame.pack(fill=tk.BOTH, expand=True, padx=15, pady=5)

        self.history_text = scrolledtext.ScrolledText(
            history_frame,
            font=(self.chinese_font, 12),
            wrap=tk.WORD,
            padx=10,
            pady=10,
            bg="#f9f9f9"
        )
        self.history_text.pack(fill=tk.BOTH, expand=True)
        self.history_text.config(state=tk.DISABLED)

        # 状态栏
        self.status_var = tk.StringVar()
        self.status_var.set("系统就绪")
        status_bar = tk.Label(
            self.root,
            textvariable=self.status_var,
            relief=tk.SUNKEN,
            anchor=tk.W,
            font=(self.chinese_font, 10),
            bg="#ecf0f1"
        )
        status_bar.pack(side=tk.BOTTOM, fill=tk.X)

        # 文本样式
        self.history_text.tag_config("user", foreground="#2c3e50", font=(self.chinese_font, 12, "bold"))
        self.history_text.tag_config("assistant", foreground="#27ae60", font=(self.chinese_font, 12))

    def ask_question(self):
        """处理用户提问"""
        question = self.question_entry.get().strip()
        if not question:
            messagebox.showwarning("提示", "请输入您的问题！")
            return

        self.display_message(f"您: {question}", "user")
        self.question_entry.delete(0, tk.END)
        self.status_var.set("正在分析问题...")
        self.root.update()

        try:
            start_time = time.time()

            # 1. 准备上下文
            context = self.prepare_context(question)
            if not context:
                self.display_message("顾问: 未找到相关内容", "assistant")
                return

            logger.debug(f"使用的上下文:\n{context[:500]}...")

            # 2. 构建提示词
            prompt = f"""你是一位专业的投资顾问，请根据以下文档内容回答问题：

【相关文档内容】
{context}

【用户问题】
{question}

请严格：
1. 基于文档内容回答
2. 引用时注明章节
3. 分点说明关键观点

请按格式回答：
💡 核心观点：
...
📖 参考章节：..."""

            # 3. 调用豆包API
            self.status_var.set("正在咨询豆包大模型...")
            self.root.update()

            logger.info("正在调用豆包API...")
            response = self.query_doubao(prompt)

            # 4. 显示结果
            self.display_message(f"顾问:\n{response}", "assistant")
            self.status_var.set(f"已回答 | 耗时: {time.time() - start_time:.1f}秒")

        except Exception as e:
            logger.error(f"提问出错: {str(e)}", exc_info=True)
            self.display_message("顾问: 处理问题时出错，请稍后再试", "assistant")
            self.status_var.set(f"错误: {str(e)}")

    def prepare_context(self, question):
        """准备上下文"""
        try:
            # 1. 使用jieba分词
            keywords = set(word for word in jieba.cut(question) if len(word) >= 2)
            logger.debug(f"分词结果: {keywords}")

            # 2. 优先搜索章节标题
            matched_sections = []
            for section in self.section_index:
                if any(keyword in section for keyword in keywords):
                    matched_sections.append(section)

            # 3. 组合相关内容
            context = ""
            for section in matched_sections[:3]:  # 最多取3个章节
                context += f"【{section}】\n{self.section_index[section][:2000]}\n\n"

            # 4. 如果没有匹配章节，使用文档开头
            return context[:MAX_CONTEXT_LENGTH] if context else self.doc_text[:MAX_CONTEXT_LENGTH]

        except Exception as e:
            logger.error(f"准备上下文出错: {str(e)}")
            return self.doc_text[:MAX_CONTEXT_LENGTH]

    def query_doubao(self, prompt):
        """调用豆包大模型"""
        for attempt in range(MAX_RETRY):
            try:
                logger.info(f"尝试第 {attempt + 1} 次API调用，提示词长度: {len(prompt)}")

                response = self.client.chat.completions.create(
                    model="deepseek-v3-250324",
                    messages=[
                        {"role": "system", "content": "你是专业的投资顾问"},
                        {"role": "user", "content": prompt},
                    ],
                    stream=True,
                    temperature=0.3,
                    max_tokens=2000,
                    timeout=30
                )

                # 处理流式响应
                full_response = ""
                for chunk in response:
                    if chunk.choices and chunk.choices[0].delta.content:
                        content = chunk.choices[0].delta.content
                        full_response += content
                        # 实时更新界面
                        self.root.update()

                logger.info("API调用成功")
                return full_response.strip()

            except Exception as e:
                logger.error(f"第 {attempt + 1} 次尝试失败: {str(e)}")
                if attempt < MAX_RETRY - 1:
                    time.sleep(RETRY_DELAY * (attempt + 1))
                else:
                    raise Exception(f"API调用失败: {str(e)}")

        return "抱歉，暂时无法获取回答。请检查网络连接后重试。"

    def display_message(self, message, tag):
        """显示消息"""
        self.history_text.config(state=tk.NORMAL)
        self.history_text.insert(tk.END, "\n\n" if tag == "user" else "\n", None)
        self.history_text.insert(tk.END, message, tag)
        self.history_text.see(tk.END)
        self.history_text.config(state=tk.DISABLED)

    def clear_conversation(self):
        """清空对话"""
        self.history_text.config(state=tk.NORMAL)
        self.history_text.delete(1.0, tk.END)
        self.history_text.config(state=tk.DISABLED)
        self.status_var.set("对话已清空")


if __name__ == "__main__":
    try:
        root = tk.Tk()
        root.option_add("*Font", "TkDefaultFont")
        app = InvestmentQAApp(root)
        root.mainloop()
    except Exception as e:
        logger.error(f"程序崩溃: {str(e)}", exc_info=True)
        messagebox.showerror("错误", f"程序启动失败:\n{str(e)}")